import os
import shutil
from random import shuffle

# 定义源目录和目标目录
dataset_dir = 'dataset'
train_dir = os.path.join(dataset_dir, 'train')
vaild_dir = os.path.join(dataset_dir, 'valid')

images_dir = 'images'
labels_dir = 'labels'

# 创建目标目录如果它们不存在
for subdir in [os.path.join(vaild_dir, images_dir), os.path.join(vaild_dir, labels_dir)]:
    if not os.path.exists(subdir):
        os.makedirs(subdir)

# 获取train目录下images的所有文件（不包括后缀）
train_images = set([os.path.splitext(f)[0] for f in os.listdir(os.path.join(train_dir, images_dir))])

# 获取train目录下labels的所有文件（不包括后缀）
train_labels = set([os.path.splitext(f)[0] for f in os.listdir(os.path.join(train_dir, labels_dir))])

# 确保图片和标签文件名（不包括后缀）是一一对应的
assert train_images == train_labels, "Images and labels do not match."

# 获取完整的图片文件列表
full_train_images = os.listdir(os.path.join(train_dir, images_dir))

# 随机打乱文件列表
shuffle(full_train_images)

# 计算要移动到vaild的文件数量（20%）
split_point = int(len(full_train_images) * 0.2)

# 分割文件列表
vaild_images = full_train_images[:split_point]
train_images = full_train_images[split_point:]

# 移动文件
for image_name in vaild_images:
    # 构建完整的文件路径
    src_image_path = os.path.join(train_dir, images_dir, image_name)
    dst_image_path = os.path.join(vaild_dir, images_dir, image_name)
    shutil.move(src_image_path, dst_image_path)

    # 构建对应的标签文件名并移动
    label_name = os.path.splitext(image_name)[0] + '.txt'
    src_label_path = os.path.join(train_dir, labels_dir, label_name)
    dst_label_path = os.path.join(vaild_dir, labels_dir, label_name)
    shutil.move(src_label_path, dst_label_path)

print("Files have been split and moved successfully.")